"""
LightRAG主类
集成文档处理、向量处理和生成等功能
"""

import os
from typing import List, Dict, Optional, Union
import logging
from pathlib import Path
import yaml
import asyncio
from dataclasses import dataclass

from .core.document import DocumentProcessor, Document
from .core.embedding import Embedding
from .core.generation import Generator, GenerationResult

logger = logging.getLogger(__name__)

@dataclass
class RAGResponse:
    """RAG响应"""
    answer: str
    sources: List[Dict]
    score: float
    followup_questions: Optional[List[str]] = None
    
class LightRAG:
    """LightRAG主类"""
    
    def __init__(self, config_path: Optional[str] = None):
        # 加载配置
        self.config = self._load_config(config_path)
        
        # 初始化组件
        self.doc_processor = DocumentProcessor(self.config.get("document", {}))
        self.embedding = Embedding(self.config.get("embedding", {}))
        self.generator = Generator(self.config.get("generation", {}))
        
        # 初始化向量存储
        self.embedding.initialize()
        
        logger.info("LightRAG初始化完成")
        
    def _load_config(self, config_path: Optional[str]) -> Dict:
        """加载配置文件"""
        if config_path is None:
            # 使用默认配置
            return {
                "document": {
                    "chunk_size": 500,
                    "chunk_overlap": 50,
                    "min_chunk_size": 100
                },
                "embedding": {
                    "model_name": "BAAI/bge-small-zh-v1.5",
                    "device": "cuda",
                    "batch_size": 32
                },
                "generation": {
                    "model_name": "gpt-3.5-turbo",
                    "temperature": 0.7,
                    "max_tokens": 1000
                }
            }
            
        # 从文件加载配置
        with open(config_path, "r", encoding="utf-8") as f:
            return yaml.safe_load(f)
            
    async def add_document(self, file_path: Union[str, Path]) -> None:
        """添加文档"""
        try:
            # 加载文档
            document = await self.doc_processor.load_document(file_path)
            
            # 处理文档
            processed_doc = self.doc_processor.process_document(document)
            
            # 添加到向量存储
            if processed_doc.chunks:
                self.embedding.add_documents(processed_doc.chunks)
                
            logger.info(f"成功添加文档: {file_path}")
            
        except Exception as e:
            logger.error(f"添加文档失败 {file_path}: {str(e)}")
            raise
            
    async def add_documents(self, file_paths: List[Union[str, Path]]) -> None:
        """批量添加文档"""
        for file_path in file_paths:
            await self.add_document(file_path)
            
    async def query(self,
                   question: str,
                   top_k: int = 3,
                   generate_followup: bool = True) -> RAGResponse:
        """查询"""
        try:
            # 搜索相关文档
            chunks = self.embedding.search(question, top_k=top_k)
            
            if not chunks:
                return RAGResponse(
                    answer="未找到相关信息",
                    sources=[],
                    score=0.0
                )
                
            # 生成回答
            result = await self.generator.generate(question, chunks)
            
            # 生成后续问题
            followup_questions = None
            if generate_followup:
                followup_questions = await self.generator.generate_followup_questions(
                    question,
                    result.answer
                )
                
            return RAGResponse(
                answer=result.answer,
                sources=result.sources,
                score=result.score,
                followup_questions=followup_questions
            )
            
        except Exception as e:
            logger.error(f"查询失败: {str(e)}")
            raise
            
    def save(self, path: str) -> None:
        """保存向量存储"""
        self.embedding.save(path)
        
    def load(self, path: str) -> None:
        """加载向量存储"""
        self.embedding.load(path)
        
    def clear(self) -> None:
        """清空向量存储"""
        self.embedding.clear() 